#include <string>
#include <iostream>
#include <fstream>
#include <math.h>
#include <random>
#include <ros/ros.h>
#include <ros/console.h>
#include <nav_msgs/Path.h>
#include <nav_msgs/Odometry.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/Point.h>
#include <visualization_msgs/MarkerArray.h>
#include <visualization_msgs/Marker.h>
#include <quadrotor_msgs/PolynomialTrajectory.h>
#include <sensor_msgs/Joy.h>
#include <algorithm>
#include <quadrotor_msgs/FlatTarget.h>
// Useful customized headers
#include "trajectory_generator_waypoint.h"
#include "sample_waypoints.h"


using namespace std;
using namespace Eigen;

// Param from launch file
double _vis_traj_width;
double _Vel, _Acc;
double _startVel_x;
int _dev_order, _min_order;

// ros related
ros::Subscriber _way_pts_sub;
ros::Publisher _wp_traj_vis_pub, _wp_path_vis_pub,pos_cmd_pub,_wp_target_vis_pub;

// for planning
int _poly_num1D;
MatrixXd _polyCoeff;
VectorXd _polyTime;
Vector3d _startPos  = Vector3d::Zero();
Vector3d _startVel = Vector3d::Zero();
Vector3d _startAcc  = Vector3d::Zero();
Vector3d _endVel = Vector3d::Zero();
Vector3d _endAcc    = Vector3d::Zero();



/* ************************************* */
bool traj_generation = false;
ros::Time start_time_;
double traj_duration_;
quadrotor_msgs::FlatTarget cmd;
VectorXd _polyTime_sum;
bool is_odom_ready;
nav_msgs::Odometry odom;
string waypoint_type = string("manual");

VectorXd getTimeSum(VectorXd time);
Vector3d getVelPoly( MatrixXd polyCoeff, int k, double t );
Vector3d getAccPoly( MatrixXd polyCoeff, int k, double t );
Vector3d getJerkPoly( MatrixXd polyCoeff, int k, double t );
void cmdCallback(const ros::TimerEvent& e);
void visWayTargetPoint( Vector3d pos,Vector3d vel);

// declare
void visWayPointTraj(MatrixXd polyCoeff, VectorXd time);
void visWayPointPath(MatrixXd path);
Vector3d getPosPoly(MatrixXd polyCoeff, int k, double t);
VectorXd timeAllocation(MatrixXd Path);
void trajGeneration(Eigen::MatrixXd path);
void rcvWaypointsCallBack(const nav_msgs::Path &wp);
void goal_callback(const geometry_msgs::PoseStamped::ConstPtr &msg);
void odom_callback(const nav_msgs::Odometry::ConstPtr &msg);

void goal_callback(const geometry_msgs::PoseStamped::ConstPtr &msg)
{
    if (!is_odom_ready) {
        ROS_ERROR("[waypoint_generator] No odom!");
        return;
    }

    ros::NodeHandle n("~");
    n.param("waypoint_type", waypoint_type, string("manual"));

    nav_msgs::Path waypoints;

    geometry_msgs::PoseStamped pt;
    pt.pose.position.x =  odom.pose.pose.position.x;
    pt.pose.position.y =  odom.pose.pose.position.y;
    pt.pose.position.z =  odom.pose.pose.position.z;

    

    if (waypoint_type == string("circle"))
    {
        waypoints.poses.push_back(pt); 
        circle(waypoints);
    }
    else if (waypoint_type == string("eight"))
    {
        waypoints = eight();
    }

    Eigen::MatrixXd waypoints_matrix(waypoints.poses.size(), 3);
    // Vector3d pt;
    for (int k = 0; k < (int)waypoints.poses.size(); k++)
    {
        waypoints_matrix.row(k) <<  waypoints.poses[k].pose.position.x , waypoints.poses[k].pose.position.y ,waypoints.poses[k].pose.position.z;
    }
    trajGeneration(waypoints_matrix);
}
void odom_callback(const nav_msgs::Odometry::ConstPtr &msg)
{
    is_odom_ready = true;
    odom = *msg;
}
//Get the path points 
void rcvWaypointsCallBack(const nav_msgs::Path & wp)
{   
    vector<Vector3d> wp_list;
    wp_list.clear();

    for (int k = 0; k < (int)wp.poses.size(); k++)
    {
        Vector3d pt( wp.poses[k].pose.position.x, wp.poses[k].pose.position.y, wp.poses[k].pose.position.z);
        wp_list.push_back(pt);
        ROS_INFO("waypoint%d: (%f, %f, %f)", k+1, pt(0), pt(1), pt(2));
    }

    MatrixXd waypoints(wp_list.size(), 3);  //add the original point
    // waypoints.row(0) = _startPos;

    for(int k = 0; k < (int)wp_list.size(); k++)
        waypoints.row(k) = wp_list[k];

    //Trajectory generation: use minimum jerk/snap trajectory generation method
    //waypoints is the result of path planning (Manual in this project)
    trajGeneration(waypoints);
}

void trajGeneration(Eigen::MatrixXd path)
{
    ros::Time time_start = ros::Time::now();

    TrajectoryGeneratorWaypoint trajectoryGeneratorWaypoint;
    
    MatrixXd vel  = MatrixXd::Zero(2, 3); 
    MatrixXd acc  = MatrixXd::Zero(2, 3);
    vel.row(0)  = _startVel;
    vel.row(1)  = _endVel;
    acc.row(0)  = _startAcc;
    acc.row(1)  = _endAcc;

    // use "trapezoidal velocity" time allocation
    _polyTime  = timeAllocation(path);

    // generate a minimum-jerk/snap piecewise monomial polynomial-based trajectory
    _polyCoeff = trajectoryGeneratorWaypoint.PolyQPGeneration(_dev_order, path, vel, acc, _polyTime);

    ros::Time time_end = ros::Time::now();
    ROS_WARN("Time consumed in trajectory generation is %f ms", (time_end - time_start).toSec() * 1000.0);

    /* ************************************* */
    traj_generation = true; 
    start_time_ = ros::Time::now();

    visWayPointPath(path);    // visulize path
    visWayPointTraj( _polyCoeff, _polyTime);    // visulize trajectory
}

int main(int argc, char** argv)
{
    ros::init(argc, argv, "traj_node");
    ros::NodeHandle nh("~");

    nh.param("planning/vel", _Vel, 1.0);
    nh.param("planning/_startVel_x", _startVel_x, 1.0);
    nh.param("planning/acc", _Acc, 1.0);
    nh.param("planning/dev_order", _dev_order, 3);  // the order of derivative, _dev_order = 3->minimum jerk, _dev_order = 4->minimum snap
    nh.param("planning/min_order", _min_order, 3);
    nh.param("vis/vis_traj_width", _vis_traj_width, 0.15);

    // set startVel endVel
    _startVel << _startVel_x , 0.0, 0.0;
    _endVel << _startVel_x , 0.0, 0.0;

    //_poly_numID is the maximum order of polynomial
    _poly_num1D = 2 * _dev_order;

    // _way_pts_sub     = nh.subscribe( "waypoints", 1, rcvWaypointsCallBack );

    _wp_traj_vis_pub = nh.advertise<visualization_msgs::Marker>("vis_trajectory", 1);
    _wp_path_vis_pub = nh.advertise<visualization_msgs::Marker>("vis_waypoint_path", 1);

    _wp_target_vis_pub = nh.advertise<visualization_msgs::Marker>("vis_waypoint_target", 1);

    /* ************************************* */
    ros::Timer cmd_timer = nh.createTimer(ros::Duration(0.01), cmdCallback);
    pos_cmd_pub = nh.advertise<quadrotor_msgs::FlatTarget>("/reference/flatsetpoint", 50);
    ros::Subscriber odom_sub = nh.subscribe("odom", 10, odom_callback);
    ros::Subscriber goal_sub = nh.subscribe("goal", 10, goal_callback);

    ros::Rate rate(100);
    bool status = ros::ok();
    while(status) 
    {
        ros::spinOnce();
        status = ros::ok();
        rate.sleep();
    }
    return 0;
}

void visWayTargetPoint( Vector3d pos,Vector3d vel)
{
    visualization_msgs::Marker _traj_vis;

    _traj_vis.header.stamp       = ros::Time::now();
    _traj_vis.header.frame_id    = "/map";

    _traj_vis.ns = "traj_node/visWayTargetPoint";
    _traj_vis.id = 0;
    _traj_vis.type = visualization_msgs::Marker::ARROW;
    _traj_vis.action = visualization_msgs::Marker::ADD;
    _traj_vis.scale.x = _vis_traj_width;
    _traj_vis.scale.y = _vis_traj_width;
    _traj_vis.scale.z = _vis_traj_width;
    _traj_vis.pose.orientation.x = 0.0;
    _traj_vis.pose.orientation.y = 0.0;
    _traj_vis.pose.orientation.z = 0.0;
    _traj_vis.pose.orientation.w = 1.0;

    _traj_vis.color.a = 1.0;
    _traj_vis.color.r = 1.0;
    _traj_vis.color.g = 0.0;
    _traj_vis.color.b = 0.0;

    double traj_len = 0.0;
    int count = 0;
    Vector3d cur, pre;
    cur.setZero();
    pre.setZero();

    _traj_vis.points.clear();
    geometry_msgs::Point pt;

    pt.x = pos(0);
    pt.y = pos(1);
    pt.z = pos(2);
    _traj_vis.points.push_back(pt);

    pt.x = pos(0) + vel(0);
    pt.y = pos(1) + vel(1);
    pt.z = pos(2) + vel(2);
    _traj_vis.points.push_back(pt);

    _wp_target_vis_pub.publish(_traj_vis);
}

void visWayPointTraj( MatrixXd polyCoeff, VectorXd time)
{
    visualization_msgs::Marker _traj_vis;

    _traj_vis.header.stamp       = ros::Time::now();
    _traj_vis.header.frame_id    = "/map";

    _traj_vis.ns = "traj_node/trajectory_waypoints";
    _traj_vis.id = 0;
    _traj_vis.type = visualization_msgs::Marker::SPHERE_LIST;
    _traj_vis.action = visualization_msgs::Marker::ADD;
    _traj_vis.scale.x = _vis_traj_width;
    _traj_vis.scale.y = _vis_traj_width;
    _traj_vis.scale.z = _vis_traj_width;
    _traj_vis.pose.orientation.x = 0.0;
    _traj_vis.pose.orientation.y = 0.0;
    _traj_vis.pose.orientation.z = 0.0;
    _traj_vis.pose.orientation.w = 1.0;

    _traj_vis.color.a = 1.0;
    _traj_vis.color.r = 1.0;
    _traj_vis.color.g = 0.0;
    _traj_vis.color.b = 0.0;

    double traj_len = 0.0;
    int count = 0;
    Vector3d cur, pre;
    cur.setZero();
    pre.setZero();

    _traj_vis.points.clear();
    Vector3d pos;
    geometry_msgs::Point pt;


    for(int i = 0; i < time.size(); i++ )   // go through each segment
    {   
        for (double t = 0.0; t < time(i); t += 0.01, count += 1)
        {
          pos = getPosPoly(polyCoeff, i, t);
          cur(0) = pt.x = pos(0);
          cur(1) = pt.y = pos(1);
          cur(2) = pt.z = pos(2);
          _traj_vis.points.push_back(pt);

          if (count) traj_len += (pre - cur).norm();
          pre = cur;
        }
    }
    ROS_WARN("Trajectory length is %f m", traj_len);
    _wp_traj_vis_pub.publish(_traj_vis);
}

void visWayPointPath(MatrixXd path)
{
    visualization_msgs::Marker points, line_list;
    int id = 0;
    points.header.frame_id    = line_list.header.frame_id    = "/map";
    points.header.stamp       = line_list.header.stamp       = ros::Time::now();
    points.ns                 = line_list.ns                 = "wp_point";
    points.action             = line_list.action             = visualization_msgs::Marker::ADD;
    points.pose.orientation.w = line_list.pose.orientation.w = 1.0;
    points.pose.orientation.x = line_list.pose.orientation.x = 0.0;
    points.pose.orientation.y = line_list.pose.orientation.y = 0.0;
    points.pose.orientation.z = line_list.pose.orientation.z = 0.0;

    points.id    = id;
    line_list.id = id;

    points.type    = visualization_msgs::Marker::SPHERE_LIST;
    line_list.type = visualization_msgs::Marker::LINE_STRIP;

    points.scale.x = 0.3;
    points.scale.y = 0.3;
    points.scale.z = 0.3;
    points.color.a = 1.0;
    points.color.r = 0.0;
    points.color.g = 0.0;
    points.color.b = 0.0;

    line_list.scale.x = 0.15;
    line_list.scale.y = 0.15;
    line_list.scale.z = 0.15;
    line_list.color.a = 1.0;

    
    line_list.color.r = 0.0;
    line_list.color.g = 1.0;
    line_list.color.b = 0.0;
    
    line_list.points.clear();

    for(int i = 0; i < path.rows(); i++){
      geometry_msgs::Point p;
      p.x = path(i, 0);
      p.y = path(i, 1); 
      p.z = path(i, 2); 

      points.points.push_back(p);
      line_list.points.push_back(p);
    }

    _wp_path_vis_pub.publish(points);
    _wp_path_vis_pub.publish(line_list);
}

Vector3d getPosPoly( MatrixXd polyCoeff, int k, double t )
{
    Vector3d ret;

    for ( int dim = 0; dim < 3; dim++ )
    {
        VectorXd coeff = (polyCoeff.row(k)).segment( dim * _poly_num1D, _poly_num1D );
        VectorXd time  = VectorXd::Zero( _poly_num1D );
        
        for(int j = 0; j < _poly_num1D; j ++)
          if(j==0)
              time(j) = 1.0;
          else
              time(j) = pow(t, j);

        ret(dim) = coeff.dot(time);
        //cout << "dim:" << dim << " coeff:" << coeff << endl;
    }
    return ret;
}

VectorXd timeAllocation( MatrixXd Path)
{ 
    VectorXd time(Path.rows() - 1);

    // The time allocation is many relative timelines but not one common timeline
    for(int i = 0; i < time.rows(); i++)
    {
        double distance = (Path.row(i+1) - Path.row(i)).norm();    // or .lpNorm<2>()
        double x1 = _Vel * _Vel / (2 * _Acc); 
        double x2 = distance - 2 * x1;
        double t1 = _Vel / _Acc;
        double t2 = x2 / _Vel;
        time(i) = 2 * t1 + t2;
    }
    cout << "timeAllocation " << time.rows() << endl;
    _polyTime_sum = getTimeSum(time);
    cout << "timeAllocation _polyTime_sum :" << _polyTime_sum.rows()  << endl;
    return time;
}


/* ************************************* */
Vector3d getVelPoly( MatrixXd polyCoeff, int k, double t )
{
    Vector3d ret = VectorXd::Zero(3);
    VectorXd _coeff(8);
    _coeff << 0,1,2,3,4,5,6,7;
    for ( int dim = 0; dim < 3; dim++ )
    {
        VectorXd coeff = (polyCoeff.row(k)).segment( dim * _poly_num1D, _poly_num1D );
        VectorXd time  = VectorXd::Zero( _poly_num1D );
        
        for(int j = 1; j < _poly_num1D; j ++)
        {
            if(j==1)
                time(j) = 1.0;
            else
                time(j) = pow(t, j-1);

            ret(dim) += _coeff(j) * time(j) * coeff(j);
        }  

        // ret(dim) = coeff.dot(time);
        //cout << "dim:" << dim << " coeff:" << coeff << endl;
    }
    return ret;
}

Vector3d getAccPoly( MatrixXd polyCoeff, int k, double t )
{
    Vector3d ret = VectorXd::Zero(3);
    VectorXd _coeff(8);
    _coeff << 0,0,2,6,12,20,30,42;
    for ( int dim = 0; dim < 3; dim++ )
    {
        VectorXd coeff = (polyCoeff.row(k)).segment( dim * _poly_num1D, _poly_num1D );
        VectorXd time  = VectorXd::Zero( _poly_num1D );
        
        for(int j = 2; j < _poly_num1D; j ++)
        {
            if(j==2)
              time(j) = 1.0;
            else
              time(j) = pow(t, j-2);
            ret(dim) += _coeff(j) * time(j) * coeff(j);
        }

        // ret(dim) = coeff.dot(time);
        //cout << "dim:" << dim << " coeff:" << coeff << endl;
    }
    return ret;
}

Vector3d getJerkPoly( MatrixXd polyCoeff, int k, double t )
{
    Vector3d ret = VectorXd::Zero(3);
    VectorXd _coeff(8);
    _coeff << 0,0,0,6,24,60,120,210;
    for ( int dim = 0; dim < 3; dim++ )
    {
        VectorXd coeff = (polyCoeff.row(k)).segment( dim * _poly_num1D, _poly_num1D );
        VectorXd time  = VectorXd::Zero( _poly_num1D );
        
        for(int j = 3; j < _poly_num1D; j ++)
        {
            if(j==3)
                time(j) = 1.0;
            else
                time(j) = pow(t, j-3);

            ret(dim) += _coeff(j) * time(j) * coeff(j);
        }
        //cout << "dim:" << dim << " coeff:" << coeff << endl;
    }
    return ret;
}
VectorXd getTimeSum(VectorXd time)
{
    VectorXd time_sum(time.rows());
    
    time_sum(0) = time(0);
    for (int i = 1; i < time.rows(); i++) {
      time_sum(i) =  time_sum(i-1) + time[i];
    }
    return time_sum;
}

void cmdCallback(const ros::TimerEvent& e) {
  /* no publishing before receive traj_ */
  
  if (!traj_generation) return;
  
  ros::Time time_now = ros::Time::now();
  double t_cur = (time_now - start_time_).toSec();
    cout << "t_cur  " << t_cur << endl;
  Eigen::Vector3d pos, vel, acc, jerk;


  int i = 0;
  double t = 0;  
  if (t_cur > _polyTime_sum(_polyTime_sum.rows()-1))
  {
    i = _polyTime_sum.rows() - 1;
    t = _polyTime_sum(i) - _polyTime_sum(i-1);
    pos = getPosPoly(_polyCoeff, i, t);
    vel.setZero();
    acc.setZero();
    jerk.setZero();
    start_time_ = ros::Time::now();
    
  }
  else if(t_cur >= 0.0)
  {
    while (true)
    {
        if (t_cur < _polyTime_sum(i))
        {
            if(i == 0)
                t = t_cur;
            else
                t = t_cur - _polyTime_sum(i-1);

            break;
        }
        i++;
    }
    pos = getPosPoly(_polyCoeff, i, t);
    vel = getVelPoly(_polyCoeff, i, t);
    acc = getAccPoly(_polyCoeff, i, t);
    jerk = getJerkPoly(_polyCoeff, i, t);
  }else {
    cout << "[Traj server]: invalid time." << endl;
  }
  cmd.header.stamp = time_now;
  cmd.header.frame_id = "world";
  cmd.type_mask = 1;


  cmd.position.x = pos(0);
  cmd.position.y = pos(1);
  cmd.position.z = pos(2);

  cmd.velocity.x = vel(0);
  cmd.velocity.y = vel(1);
  cmd.velocity.z = vel(2);

  cmd.acceleration.x = acc(0);
  cmd.acceleration.y = acc(1);
  cmd.acceleration.z = acc(2);

  cmd.jerk.x = jerk(0);
  cmd.jerk.y = jerk(1);
  cmd.jerk.z = jerk(2);
  pos_cmd_pub.publish(cmd);

    visWayTargetPoint(pos,vel);

}